import os
import sys
from time import sleep
from netaddr import *
from nmb.NetBIOS import NetBIOS
from smb.SMBConnection import SMBConnection
import fnmatch
import json
""" 连接 """
def connect(host, port, username, password):
    try:
        net = NetBIOS()
        net_name = str(net.queryIPForName(host)[0])
        net.close()
        conn = SMBConnection(username, password, 'smb2',net_name, use_ntlm_v2=False)
        if conn.connect(host, port, timeout=3):
            return conn
        else:
            raise ValueError("Connection error2: %s" % (host))
    except:
        raise ValueError("Connection error1: %s" % (host))
""" 搜索单文件 """
def find_file(conn, share_name, remote_path, keyword):
    files = conn.listPath(share_name, remote_path)
    for file in files:
        try:
            if file.filename != "." and file.filename != "..":
                if file.isDirectory:
                    # 递归调用，搜索子目录
                    res = find_file(conn, share_name, f"{remote_path}/{file.filename}", keyword)
                    if res != "":
                        return res
                else:
                    # 检查文件名是否包含关键词
                    if fnmatch.fnmatch(file.filename, keyword):
                        # 构建完整路径并添加到列表中
                        full_path = f"{remote_path}/{file.filename}"
                        return full_path
                        
        except Exception as e:
            raise ValueError(f"搜索文件时出错: {remote_path}. 错误: {e}")
    return ""
""" 搜索所有文件 """
def find_files(conn, share_name, remote_path, keyword):
    res = []
    files = conn.listPath(share_name, remote_path)
    for file in files:
        try:
            if file.filename != "." and file.filename != "..":
                if file.isDirectory:
                    # 递归调用，搜索子目录
                    res.extend(find_files(conn, share_name, f"{remote_path}/{file.filename}", keyword))
                else:
                    # 检查文件名是否包含关键词
                    if fnmatch.fnmatch(file.filename, key):
                        # 构建完整路径并添加到列表中
                        full_path = f"{remote_path}/{file.filename}"
                        res.append(f"{full_path}")
        except Exception as e:
        # raise ValueError(f"搜索文件时出错: {remote_path}. 错误: {e}")
            pass
    return res
""" 搜索所有目录 """
def find_dirs(conn, share_name, remote_path, keyword):
    res = []
    files = conn.listPath(share_name, remote_path)
    for file in files:
        try:
            if file.filename != "." and file.filename != "..":
                if file.isDirectory:
                    # 检查文件名是否包含关键词
                    if fnmatch.fnmatch(file.filename, keyword):
                        # 构建完整路径并添加到列表中
                        full_path = f"{remote_path}/{file.filename}"
                        res.append(full_path)
                    else:
                        # 递归调用，搜索子目录
                        res.extend(find_dirs(conn, share_name, f"{remote_path}/{file.filename}", keyword))
        except Exception as e:
            # raise ValueError(f"搜索文件时出错: {remote_path}. 错误: {e}")
            pass
    return res
""" 搜索所有目录 """
def find_dir(conn, share_name, remote_path, keyword):
    files = conn.listPath(share_name, remote_path)
    for file in files:
        try:
            if file.filename != "." and file.filename != "..":
                if file.isDirectory:
                    # 检查文件名是否包含关键词
                    if fnmatch.fnmatch(file.filename, keyword):
                        # 构建完整路径并添加到列表中
                        full_path = f"{remote_path}/{file.filename}"
                        return full_path
                    else:
                        # 递归调用，搜索子目录
                        res = find_dir(conn, share_name, f"{remote_path}/{file.filename}", keyword)
                        if res != "":
                            return res
        except Exception as e:
            raise ValueError(f"搜索文件时出错: {remote_path}. 错误: {e}")
    return ""
""" 列表获取 """
def list(conn, share_name, path2):
    try:
        # 列出件夹中的文件和目录
        file_list = conn.listPath(share_name,path2)
        res = []
        for file in file_list:
            res.append(file.filename)
        return res
    except:
        raise ValueError("List: %s %s" % (share_name, path2))
    
# """ share列表获取 """
def list_share(conn, timeout):
    try:
        # 列出共享文件夹中的文件和目录
        share_list = conn.listShares(timeout)
        res = []
        for share in share_list:
            res.append(share.name)
        return res
    except:
        raise ValueError("List share: %d" % (timeout))
# 上传目录
def upload_directory(conn, share_name, remote_path, local_path):
    if not os.path.exists(local_path):
        raise ValueError("Local dir: %s does not exist" % (local_path))
    res = []
    # 创建远程路径（如果需要）
    if not conn.listPath(share_name, remote_path).isDirectory:
        conn.createDirectory(share_name, remote_path)
    
    try:
        files = os.listdir(local_path)
        for file in files:
            local_file_path = os.path.join(local_path, file)
            remote_file_path = os.path.join(remote_path, file)
            if os.path.isdir(local_file_path):
                res.extend(upload_directory(conn, share_name, remote_file_path, local_file_path))
            else:
                res.append(upload_file(conn, share_name, remote_file_path, local_file_path))
    except:
        raise ValueError("Upload dir: %s" % (remote_path))

""" 上传 """
def upload_file(conn,share_name, local_path, remote_path):
    try:
        with open(local_path, 'rb') as local_file:
            conn.storeFile(share_name, remote_path, local_file)
        return remote_path
    except:
        raise ValueError("Upload file remote: %s" % (remote_path))
""" # 下载目录 """
def download_directory(conn, share_name, remote_path, local_path):
    res = []
    if not os.path.exists(local_path):
        os.makedirs(local_path)
    try:
        files = conn.listPath(share_name, remote_path)
        for file in files:
            if file.filename != "." and file.filename != "..":
                if file.isDirectory:
                    res.extend(download_directory(conn, share_name, f"{remote_path}/{file.filename}", f"{local_path}/{file.filename}"))
                else:
                    res.append(download_file(conn, share_name, f"{remote_path}/{file.filename}", f"{local_path}/{file.filename}"))
    except:
        raise ValueError("Download dir: %s" % (remote_path))
    return res
""" # 下载文件 """
def download_file(conn, share_name, remote_path, local_path):
    try:
        if remote_path != "." and remote_path != "..":
            with open(local_path, 'wb') as local_file:
                conn.retrieveFile(share_name, remote_path, local_file)
            return local_path
    except:
        raise ValueError("Download file: %s" % (remote_path))

if __name__ == "__main__":
    msg = "Usage: smb2.exe <Task> <Host> <Port> <Username> <Password> <ShareName>\n\n <Task>[\nlist@/test/test2 \nlist_share@3 \nupload_file@D:\\src.txt;/test/dst.txt \nupload_dir@D:\\test;/test/test2 \ndownload_file@/test/src.txt;D:\\test.txt \ndownload_dir@/test;D:\\test \n]"
    try:
        if len(sys.argv) == 7:
            task = sys.argv[1]
            host = sys.argv[2]
            port = int(sys.argv[3])
            username = sys.argv[4]
            password = sys.argv[5]
            share_name = sys.argv[6]
            if "list@" in task:
                path = task.replace("list@","")
                conn = connect(host, port, username, password)
                data = list(conn, share_name, path)
                print(json.dumps(data))
                conn.close()
            elif "list_share@" in task:
                timeout = int(task.replace("list_share@",""))
                conn = connect(host, port, username, password)
                data = list_share(conn, timeout)
                print(json.dumps(data))
                conn.close()
            elif "upload_file@" in task:
                try:
                    res = task.replace("upload_file@","").split(";")
                    src = res[0]
                    dst = res[1]
                    conn = connect(host, port, username, password)
                    res = upload_file(conn, share_name,src,dst)
                    print(res)
                    conn.close()
                except:
                    raise ValueError("Task error: %s" % (task))
            elif "upload_dir@" in task:
                try:
                    res = task.replace("upload_dir@","").split(";")
                    src = res[0]
                    dst = res[1]
                    conn = connect(host, port, username, password)
                    data = upload_directory(conn, share_name,src,dst)
                    print(json.dumps(data))
                    conn.close()
                except:
                    raise ValueError("Task error: %s" % (task))
            elif "download_file@" in task:
                try:
                    res = task.replace("download_file@","").split(";")
                    src = res[0]
                    dst = res[1]
                    conn = connect(host, port, username, password)
                    res = download_file(conn, share_name,src,dst)
                    print(res)
                    conn.close()
                except:
                    raise ValueError("Task error: %s" % (task))
            elif "download_dir@" in task:
                try:
                    res = task.replace("download_dir@","").split(";")
                    src = res[0]
                    dst = res[1]
                    conn = connect(host, port, username, password)
                    data = download_directory(conn, share_name,src,dst)
                    print(json.dumps(data))
                    conn.close()
                except:
                    raise ValueError("Task error: %s" % (task))
            elif "find_file@" in task:
                try:
                    res = task.replace("find_file@","").split(";")
                    path = res[0]
                    key = res[1]
                    conn = connect(host, port, username, password)
                    data = find_file(conn, share_name,path,key)
                    print(data)
                    conn.close()
                except:
                    raise ValueError("Task error: %s" % (task))
            elif "find_files@" in task:
                try:
                    res = task.replace("find_files@","").split(";")
                    path = res[0]
                    key = res[1]
                    conn = connect(host, port, username, password)
                    data = find_files(conn, share_name,path,key)
                    print(json.dumps(data))
                    conn.close()
                except:
                    raise ValueError("Task error: %s" % (task))
            elif "find_dirs@" in task:
                try:
                    res = task.replace("find_files@","").split(";")
                    path = res[0]
                    key = res[1]
                    conn = connect(host, port, username, password)
                    data = find_dirs(conn, share_name,path,key)
                    print(json.dumps(data))
                    conn.close()
                except:
                    raise ValueError("Task error: %s" % (task))
            elif "find_dir@" in task:
                try:
                    res = task.replace("find_files@","").split(";")
                    path = res[0]
                    key = res[1]
                    conn = connect(host, port, username, password)
                    data = find_dir(conn, share_name,path,key)
                    print(data)
                    conn.close()
                except:
                    raise ValueError("Task error: %s" % (task))
            else:
                raise ValueError("Task error: %s" % (task))
        else:
            raise ValueError(msg)
    except:
        raise ValueError(msg)
